#pragma once

#include "MemoryTracker.h"

#include <map>
#include <iostream>
#include <string>
#include <cassert>
#include <iostream>

using namespace MemMgr;

MemoryTracker::MemoryTracker(std::ostream& oStream) : mAllocatedMemory(0), mAllocationMap(), mOutput(oStream)
{
}

MemoryTracker::~MemoryTracker()
{
	LogLeaks();
}

MemoryTracker& MemoryTracker::GetSingleton()
{ 
	static MemoryTracker instance;
	return instance; 
}

MemoryTracker* MemoryTracker::GetSingletonPtr()
{
	static MemoryTracker* instancePtr = &GetSingleton();
	return instancePtr; 
}

void MemoryTracker::RegisterAlloc(const void * const iPtr, const size_t iSize, const char* iFilename, const char* iFunctionName, const unsigned int iRow, const std::string iCategory)
{
	if(nullptr != iPtr)
	{
		AllocationInfoMap::const_iterator it = mAllocationMap.find(iPtr);

		assert(it == mAllocationMap.end());

		mAllocationMap[iPtr] = AllocationInfo(iSize, iFilename, iFunctionName, iRow, iCategory);
		mAllocatedMemory += iSize;

		//assert(it == mAllocationMap.begin() || (--it, static_cast<const char*>(iPtr) - (*it).second.size > (*it).first));
		//assert(it != mAllocationMap.end() || static_cast<const char*>(iPtr) + iSize < (*it).first);
	}
}

void MemoryTracker::UnregisterAlloc(const void * const iPtr, const size_t iSize)
{
	if(nullptr != iPtr)
	{
		AllocationInfoMap::const_iterator it = mAllocationMap.find(iPtr);

		if(it != mAllocationMap.end() && iSize == it->second.size)
		{
			mAllocatedMemory -= (*it).second.size;
			mAllocationMap.erase(it);
		} else
		{
			LogError(iPtr, iSize);
		}
	}
}

void MemoryTracker::UnregisterAlloc(const void * const iPtr)
{
	if(nullptr != iPtr)
	{
		AllocationInfoMap::const_iterator it = mAllocationMap.find(iPtr);

		if(it != mAllocationMap.end())
		{
			mAllocatedMemory -= (*it).second.size;
			mAllocationMap.erase(it);
		} else
		{
			LogError(iPtr);
		}
	}
}

void MemoryTracker::LogError(const void * const iPtr, const size_t iSize)
{
	AllocationInfoMap::const_iterator it = mAllocationMap.find(iPtr);

	if(it != mAllocationMap.end())
	{
		if	((*it).first == iPtr && iSize > (*it).second.size)
		{
			mOutput << "ERROR: Attempt to deallocate a block greater than the one allocated.\n"
				<< "Deallocation: \n"
				<< " Size: " << iSize << '\n'
				<< "Block: \n"
				<< " Category: " << (*it).second.category << '\n'
				<< " Address: " << (*it).first << '\n'
				<< " Size: " << (*it).second.size << '\n'
				<< " File: " << (*it).second.filename << '\n'
				<< " Function: " << (*it).second.functionName << '\n'
				<< " Line: " << (*it).second.line << '\n' << std::endl;
		} else
			if	(((*it).first == iPtr && iSize < (*it).second.size) ||
				(((++it)) != mAllocationMap.end() && static_cast<const char*>((*it).first) + (*it).second.size > iPtr)  ||
				((--it) != mAllocationMap.begin() && (--it, static_cast<const char*>((*it).first) + (*it).second.size > iPtr)) )
			{
				mOutput << "ERROR: Attempt to deallocate a section of a previously allocated block.\n"
					<< "Deallocation: \n"
					<< " Address: " << iPtr << '\n'
					<< " Size: " << iSize << '\n'
					<< "Containing Block:\n"
					<< " Category: " << (*it).second.category << '\n'
					<< " Address: " << (*it).first << '\n'
					<< " Size: " << (*it).second.size << '\n'
					<< " File: " << (*it).second.filename << '\n'
					<< " Function: " << (*it).second.functionName << '\n'
					<< " Line: " << (*it).second.line << '\n' << std::endl;
			}
	} else
	{
		mOutput << "ERROR: Attempt to deallocate a block never allocated.\n"
			<< "Deallocation: \n"
			<< " Address: " << iPtr << '\n'
			<< " Size: " << iSize << '\n' << std::endl;
	}
}

void MemoryTracker::LogError(const void * const iPtr)
{
	mOutput << "ERROR: Attempt to deallocate a block never allocated.\n"
		<< "Deallocation: \n"
		<< " Address: " << iPtr << '\n' << std::endl;
};

void MemoryTracker::LogLeaks()
{
	if(!mAllocationMap.empty())
	{
		mOutput << "\n\nMEMORY LEAKS DETECTED:" << '\n' << std::endl;
		mOutput << "UNALLOCATED BYTES: " << mAllocatedMemory << '\n' << std::endl;

		unsigned int i = 0;
		for(AllocationInfoMap::const_iterator it = mAllocationMap.begin(); it != mAllocationMap.end(); ++it)
		{
			mOutput << "LEAK #" << i++ << '\n'
				<< " Category: " << (*it).second.category << '\n'
				<< " Address: " << (*it).first << '\n'
				<< " Size: " << (*it).second.size << '\n'
				<< " File:  " << (*it).second.filename << '\n'
				<< " Function: "<< (*it).second.functionName << '\n'
				<< " Line: "<< (*it).second.line << '\n' << std::endl;
		}
	} else
	{
		mOutput << "NO MEMORY LEAKS DETECTED. :)" << std::endl;
	}
}